import datetime
import logging

import matplotlib.pyplot as plt
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import ConfusionMatrixDisplay
from sklearn.metrics import accuracy_score, confusion_matrix
from sklearn.model_selection import train_test_split, GridSearchCV

# 配置日志
logging.basicConfig(filename=f'./logs/random_forest.log',
                    level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    encoding='utf-8')

# 加载数据集
# dataset_path = './data/car_1000.txt'
dataset_path = './data/car_1000_numeric.csv'
logging.info(f"正在加载数据集: {dataset_path}")

data = pd.read_csv(dataset_path, header=None)

# 分离特征和标签
X = data.iloc[:, :-1]  # 特征
y = data.iloc[:, -1]  # 标签

# 处理类别变量（OneHot编码）
# X_encoded = pd.get_dummies(X)

# 将数据集划分为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建随机森林分类器
rf_clf = RandomForestClassifier()

# 设置超参数网格
param_grid = {
    'n_estimators': [50, 100, 200],
    'max_depth': [None, 10, 20, 30],
    'min_samples_split': [2, 5, 10],
    'min_samples_leaf': [1, 2, 4]
}

# 使用GridSearchCV进行超参数调优
grid_search = GridSearchCV(rf_clf, param_grid, cv=5, scoring='accuracy')
grid_search.fit(X_train, y_train)

# 输出最佳超参数组合
best_params = grid_search.best_params_
logging.info(f"最佳超参数组合: {best_params}")
print(f"最佳超参数组合: {best_params}")

# 在测试集上进行预测
y_pred = grid_search.predict(X_test)

# 计算准确率
accuracy = accuracy_score(y_test, y_pred)
logging.info(f"准确率：{accuracy}")
print(f"准确率: {accuracy}")

# 绘制混淆矩阵
conf_matrix = confusion_matrix(y_test, y_pred)
ConfusionMatrixDisplay.from_estimator(grid_search, X_test, y_test, cmap=plt.get_cmap('Blues'))

logging.info("正在保存混淆矩阵图...")
plt.savefig(f'./images/{datetime.datetime.now().strftime("%Y-%m-%d_%H_%M_%S")}_rf_confusion_matrix.png')
logging.info("混淆矩阵图已保存...")

plt.show()
